from gmssl.sm4 import CryptSM4,SM4_ENCRYPT,SM4_DECRYPT
import binascii



crypt_sm4=CryptSM4()

def tobyte(st):
    strs=st.encode('utf-8')
##    str_s=binascii.b2a_hex(strs)
##    strs=binascii.b2a_base64(strs)
    str_s=binascii.hexlify(strs)
    return str_s

def tost(byt):
    s=binascii.unhexlify(byt).decode('utf-8')
    return s


def key16(st):
    while len(st)%16!=0:
        st+=b'v'
    s1=st[:-6]
    e1=st[-6:]
    st=e1+s1
    return st

def sm4_encrypt(k,iv,st):
    crypt_sm4.set_key(k,SM4_ENCRYPT)
    enval=crypt_sm4.crypt_cbc(iv,st)
    return binascii.b2a_hex(enval).decode('utf-8')

def sm4_decrypt(k,iv,enst):
    enst=binascii.a2b_hex(enst)
    crypt_sm4.set_key(k,SM4_DECRYPT)
    deval=crypt_sm4.crypt_cbc(iv,enst)
    return deval
    

if __name__=='__main__':

    iv='abcdefg'
    choose=input('SM4加解密(en/de)：')
    if choose=='en':
        key=input('密钥：')
        schar=input('要加密的字符：')
        ##    print(key16(tobyte(iv))[:16])

        enst=sm4_encrypt(key16(tobyte(key)),key16(tobyte(iv)),tobyte(schar))
        print('加密：{}'.format(enst))
    if choose=='de':
        key=input('密钥：')
        schar=input('要解密的字符：')
        deb=sm4_decrypt(key16(tobyte(key)),key16(tobyte(iv)),schar.encode('utf-8'))
        print('解密：{}'.format(tost(deb)))
    
